package com.example.hook.runtime;

import net.bytebuddy.agent.ByteBuddyAgent;
import javassist.*;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
import java.security.ProtectionDomain;
import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 运行时Hook管理器 - 支持在程序运行过程中动态Hook方法
 */
public class RuntimeHookManager {
    
    private static Instrumentation instrumentation;
    private static final Map<String, HookInfo> activeHooks = new ConcurrentHashMap<>();
    private static final Map<String, MethodInterceptor> interceptors = new ConcurrentHashMap<>();
    
    /**
     * 初始化运行时Hook管理器
     */
    public static void init() {
        if (instrumentation == null) {
            // 使用ByteBuddy Agent获取Instrumentation实例
            instrumentation = ByteBuddyAgent.install();
            System.out.println("[RuntimeHook] 运行时Hook管理器已初始化");
        }
    }
    
    /**
     * 动态Hook指定方法
     * 
     * @param className 目标类名
     * @param methodName 目标方法名
     * @param interceptor 方法拦截器
     */
    public static boolean hookMethod(String className, String methodName, MethodInterceptor interceptor) {
        if (instrumentation == null) {
            System.err.println("[RuntimeHook] 错误: 管理器未初始化，请先调用init()");
            return false;
        }
        
        try {
            Class<?> targetClass = Class.forName(className);
            String hookKey = className + "." + methodName;
            
            // 检查是否已经Hook
            if (activeHooks.containsKey(hookKey)) {
                System.out.println("[RuntimeHook] 警告: 方法已被Hook，将替换现有Hook");
                unhookMethod(className, methodName);
            }
            
            // 创建运行时转换器
            RuntimeMethodTransformer transformer = new RuntimeMethodTransformer(className, methodName);
            
            // 注册拦截器
            interceptors.put(hookKey, interceptor);
            
            // 添加转换器并重转换类
            instrumentation.addTransformer(transformer, true);
            
            // 保存Hook信息
            HookInfo hookInfo = new HookInfo(transformer, interceptor, targetClass);
            activeHooks.put(hookKey, hookInfo);
            
            // 重转换类使Hook生效
            instrumentation.retransformClasses(targetClass);
            
            System.out.println("[RuntimeHook] 成功Hook方法: " + hookKey);
            return true;
            
        } catch (Exception e) {
            System.err.println("[RuntimeHook] Hook失败: " + className + "." + methodName);
            e.printStackTrace();
            return false;
        }
    }
    
    /**
     * 取消Hook指定方法
     */
    public static boolean unhookMethod(String className, String methodName) {
        String hookKey = className + "." + methodName;
        HookInfo hookInfo = activeHooks.remove(hookKey);
        
        if (hookInfo != null) {
            try {
                // 移除转换器
                instrumentation.removeTransformer(hookInfo.transformer);
                
                // 移除拦截器
                interceptors.remove(hookKey);
                
                // 重转换类恢复原状
                instrumentation.retransformClasses(hookInfo.targetClass);
                
                System.out.println("[RuntimeHook] 成功取消Hook: " + hookKey);
                return true;
                
            } catch (Exception e) {
                System.err.println("[RuntimeHook] 取消Hook失败: " + hookKey);
                e.printStackTrace();
                return false;
            }
        } else {
            System.out.println("[RuntimeHook] 警告: 方法未被Hook: " + hookKey);
            return false;
        }
    }
    
    /**
     * 获取当前所有活跃的Hook
     */
    public static String[] getActiveHooks() {
        return activeHooks.keySet().toArray(new String[0]);
    }
    
    /**
     * 清除所有Hook
     */
    public static void clearAllHooks() {
        String[] hooks = getActiveHooks();
        for (String hook : hooks) {
            String[] parts = hook.split("\\.");
            if (parts.length >= 2) {
                String className = hook.substring(0, hook.lastIndexOf('.'));
                String methodName = parts[parts.length - 1];
                unhookMethod(className, methodName);
            }
        }
        System.out.println("[RuntimeHook] 已清除所有Hook");
    }
    
    /**
     * 内部调用入口 - 由被Hook的方法调用
     */
    public static Object processMethodCall(String className, String methodName, 
                                         Object instance, Object[] args, 
                                         CallableMethod originalMethod) throws Exception {
        
        String hookKey = className + "." + methodName;
        MethodInterceptor interceptor = interceptors.get(hookKey);
        
        if (interceptor != null) {
            // 调用原始方法获取结果
            Object originalResult = originalMethod.call();
            
            // 通过拦截器处理
            return interceptor.intercept(instance, methodName, args, originalResult);
        } else {
            // 没有拦截器，直接调用原方法
            return originalMethod.call();
        }
    }
    
    /**
     * Hook信息存储类
     */
    private static class HookInfo {
        final ClassFileTransformer transformer;
        final MethodInterceptor interceptor;
        final Class<?> targetClass;
        
        HookInfo(ClassFileTransformer transformer, MethodInterceptor interceptor, Class<?> targetClass) {
            this.transformer = transformer;
            this.interceptor = interceptor;
            this.targetClass = targetClass;
        }
    }
    
    /**
     * 运行时方法转换器
     */
    private static class RuntimeMethodTransformer implements ClassFileTransformer {
        private final String targetClassName;
        private final String targetMethodName;
        
        public RuntimeMethodTransformer(String className, String methodName) {
            this.targetClassName = className;
            this.targetMethodName = methodName;
        }
        
        @Override
        public byte[] transform(ClassLoader loader, String className,
                              Class<?> classBeingRedefined,
                              ProtectionDomain protectionDomain,
                              byte[] classfileBuffer) throws IllegalClassFormatException {
            
            // 只处理目标类的重转换
            if (classBeingRedefined != null && 
                targetClassName.equals(classBeingRedefined.getName())) {
                
                try {
                    return modifyClassBytecode(classfileBuffer);
                } catch (Exception e) {
                    System.err.println("[RuntimeHook] 字节码转换失败: " + e.getMessage());
                    e.printStackTrace();
                    return null;
                }
            }
            
            return null; // 不修改其他类
        }
        
        private byte[] modifyClassBytecode(byte[] originalBytecode) throws Exception {
            ClassPool cp = ClassPool.getDefault();
            CtClass cc = cp.makeClass(new ByteArrayInputStream(originalBytecode));
            
            try {
                // 获取目标方法
                CtMethod[] methods = cc.getDeclaredMethods();
                CtMethod targetMethod = null;
                
                for (CtMethod method : methods) {
                    if (method.getName().equals(targetMethodName) && 
                        !method.getName().endsWith("$original")) {
                        targetMethod = method;
                        break;
                    }
                }
                
                if (targetMethod == null) {
                    System.err.println("[RuntimeHook] 未找到目标方法: " + targetMethodName);
                    return null;
                }
                
                // 检查是否已经被Hook过
                String originalMethodName = targetMethodName + "$original";
                try {
                    cc.getDeclaredMethod(originalMethodName);
                    // 如果已存在$original方法，说明已经被Hook过，直接返回
                    System.out.println("[RuntimeHook] 方法已被Hook，跳过重复处理");
                    return null;
                } catch (NotFoundException e) {
                    // 方法不存在，继续处理
                }
                
                // 创建原始方法的副本
                CtMethod originalMethod = CtNewMethod.copy(targetMethod, originalMethodName, cc, null);
                originalMethod.setModifiers(targetMethod.getModifiers() & ~Modifier.PUBLIC & ~Modifier.PROTECTED | Modifier.PRIVATE);
                cc.addMethod(originalMethod);
                
                // 修改原方法，使其调用Hook管理器
                String newBody = generateHookMethodBody(targetClassName, targetMethodName, originalMethodName, targetMethod);
                targetMethod.setBody(newBody);
                
                System.out.println("[RuntimeHook] 字节码修改完成: " + targetClassName + "." + targetMethodName);
                
                byte[] modifiedBytecode = cc.toBytecode();
                return modifiedBytecode;
                
            } finally {
                cc.detach();
            }
        }
        
        private String generateHookMethodBody(String className, String methodName, 
                                            String originalMethodName, CtMethod method) throws NotFoundException {
            
            StringBuilder body = new StringBuilder();
            body.append("{\n");
            
            // 准备参数数组
            body.append("    Object[] args = $args;\n");
            
            // 创建原始方法调用的Callable
            body.append("    com.example.hook.runtime.CallableMethod originalCall = new com.example.hook.runtime.CallableMethod() {\n");
            body.append("        public Object call() throws Exception {\n");
            
            // 根据返回类型生成不同的调用代码
            CtClass returnType = method.getReturnType();
            if (returnType == CtClass.voidType) {
                body.append("            ").append(originalMethodName).append("($$);\n");
                body.append("            return null;\n");
            } else {
                body.append("            return ($w) ").append(originalMethodName).append("($$);\n");
            }
            
            body.append("        }\n");
            body.append("    };\n");
            
            // 调用Hook管理器
            body.append("    Object result = com.example.hook.runtime.RuntimeHookManager.processMethodCall(\n");
            body.append("        \"").append(className).append("\",\n");
            body.append("        \"").append(methodName).append("\",\n");
            body.append("        this,\n");
            body.append("        args,\n");
            body.append("        originalCall\n");
            body.append("    );\n");
            
            // 根据返回类型生成返回语句
            if (returnType == CtClass.voidType) {
                body.append("    return;\n");
            } else {
                body.append("    return (").append(returnType.getName()).append(") result;\n");
            }
            
            body.append("}\n");
            
            return body.toString();
        }
    }
}